from vector_database import database

def create_prompt_from_query(question: str, template_path: str = "prompt_template.md", limit: int = 5) -> str:
    from jinja2 import Environment, FileSystemLoader

    env = Environment(loader=FileSystemLoader("templates"))
    template = env.get_template(template_path)
    
    results = []
    for r in database.query_db(question, limit):
        # get chapters list
        headings = []
        for c in r["entity"]["chapters"]:
            headings.append(r["entity"]["chapters"][c])
        
        # create record
        record = {}
        record["chapters"] = " / ".join(headings)
        record["text"] = r["entity"]["text"]

        results.append(record)
    
    # render prompt
    prompt = template.render(question=question, chunks=results)
    
    return prompt


if __name__ == "__main__":
    question = input("Enter your question: ")
    if not question:
        exit(0)
    limit = int(input("Enter the search limit: "))

    prompt = create_prompt_from_query(question, limit=limit)
    
    with open("prompt.md", mode="w") as f:
        f.write(prompt)
        print("The generated prompt is added to 'prompt.md'")
